// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/array/array_base.h>

#include <cstdint>
#include <memory>
#include <sstream>  // IWYU pragma: keep
#include <string>
#include <type_traits>
#include <utility>

#include <nebula/array/array_binary.h>
#include <nebula/array/array_dict.h>
#include <nebula/array/array_nested.h>
#include <nebula/array/array_primitive.h>
#include <nebula/array/util.h>
#include <nebula/array/validate.h>
#include <nebula/core/buffer.h>
#include <nebula/core/compare.h>
#include <nebula/core/pretty_print.h>
#include <nebula/types/scalar.h>
#include <turbo/utility/status.h>
#include <nebula/types/type.h>
#include <nebula/types/type_fwd.h>
#include <nebula/types/type_traits.h>
#include <turbo/log/logging.h>
#include <nebula/util/ree_util.h>
#include <nebula/core/visit_array_inline.h>
#include <nebula/core/visitor.h>

namespace nebula {

    class ExtensionArray;

    // ----------------------------------------------------------------------
    // Base array class

    int64_t Array::null_count() const { return data_->get_null_count(); }

    int64_t Array::compute_logical_null_count() const {
        return data_->compute_logical_null_count();
    }

    namespace internal {

        struct ScalarFromArraySlotImpl {
            template<typename T>
            using ScalarType = typename TypeTraits<T>::ScalarType;

            turbo::Status Visit(const NullArray &a) {
                out_ = std::make_shared<NullScalar>();
                return turbo::OkStatus();
            }

            turbo::Status Visit(const BooleanArray &a) { return finish(a.value(index_)); }

            template<typename T>
            turbo::Status Visit(const NumericArray<T> &a) {
                return finish(a.value(index_));
            }

            turbo::Status Visit(const Decimal128Array &a) {
                return finish(Decimal128(a.get_value(index_)));
            }

            turbo::Status Visit(const Decimal256Array &a) {
                return finish(Decimal256(a.get_value(index_)));
            }

            template<typename T>
            turbo::Status Visit(const BaseBinaryArray<T> &a) {
                return finish(a.get_string(index_));
            }

            turbo::Status Visit(const BinaryViewArray &a) { return finish(a.get_string(index_)); }

            turbo::Status Visit(const FixedSizeBinaryArray &a) { return finish(a.get_string(index_)); }

            turbo::Status Visit(const DayTimeIntervalArray &a) { return finish(a.value(index_)); }

            turbo::Status Visit(const MonthDayNanoIntervalArray &a) { return finish(a.value(index_)); }

            template<typename T>
            turbo::Status Visit(const VarLengthListLikeArray<T> &a) {
                return finish(a.value_slice(index_));
            }

            turbo::Status Visit(const FixedSizeListArray &a) { return finish(a.value_slice(index_)); }

            turbo::Status Visit(const StructArray &a) {
                ScalarVector children;
                for (const auto &child: a.fields()) {
                    children.emplace_back();
                    TURBO_MOVE_OR_RAISE(children.back(), child->get_scalar(index_));
                }
                return finish(std::move(children));
            }

            turbo::Status Visit(const SparseUnionArray &a) {
                int8_t type_code = a.type_code(index_);

                ScalarVector children;
                for (int i = 0; i < a.type()->num_fields(); ++i) {
                    children.emplace_back();
                    TURBO_MOVE_OR_RAISE(children.back(), a.field(i)->get_scalar(index_));
                }

                out_ = std::make_shared<SparseUnionScalar>(std::move(children), type_code, a.type());
                return turbo::OkStatus();
            }

            turbo::Status Visit(const DenseUnionArray &a) {
                const auto type_code = a.type_code(index_);
                // child array which stores the actual value
                auto arr = a.field(a.child_id(index_));
                // need to look up the value based on offsets
                auto offset = a.value_offset(index_);
                TURBO_MOVE_OR_RAISE(auto value, arr->get_scalar(offset));
                out_ = std::make_shared<DenseUnionScalar>(value, type_code, a.type());
                return turbo::OkStatus();
            }

            turbo::Status Visit(const DictionaryArray &a) {
                auto ty = a.type();

                TURBO_MOVE_OR_RAISE(
                        auto index, MakeScalar(turbo::checked_cast<const DictionaryType &>(*ty).index_type(),
                                               a.get_value_index(index_)));

                auto scalar = DictionaryScalar(ty);
                scalar.is_valid = a.is_valid(index_);
                scalar.value.index = index;
                scalar.value.dictionary = a.dictionary();

                out_ = std::make_shared<DictionaryScalar>(std::move(scalar));
                return turbo::OkStatus();
            }

            turbo::Status Visit(const RunEndEncodedArray &a) {
                ArraySpan span{*a.data()};
                const int64_t physical_index = ree_util::FindPhysicalIndex(span, index_, span.offset);
                ScalarFromArraySlotImpl scalar_from_values(*a.values(), physical_index);
                TURBO_MOVE_OR_RAISE(auto value, std::move(scalar_from_values).finish());
                out_ = std::make_shared<RunEndEncodedScalar>(std::move(value), a.type());
                return turbo::OkStatus();
            }

            turbo::Status Visit(const ExtensionArray &a) {
                TURBO_MOVE_OR_RAISE(auto storage, a.storage()->get_scalar(index_));
                out_ = std::make_shared<ExtensionScalar>(std::move(storage), a.type());
                return turbo::OkStatus();
            }

            template<typename Arg>
            turbo::Status finish(Arg &&arg) {
                return MakeScalar(array_.type(), std::forward<Arg>(arg)).try_value(&out_);
            }

            turbo::Status finish(std::string arg) {
                return MakeScalar(array_.type(), Buffer::from_string(std::move(arg))).try_value(&out_);
            }

            turbo::Result<std::shared_ptr<Scalar>> finish() &&{
                if (index_ >= array_.length()) {
                    return turbo::out_of_range_error("index with value of ", index_,
                                              " is out-of-bounds for array of length ",
                                              array_.length());
                }

                // Skip checking for nulls in RUN_END_ENCODED arrays to avoid potentially
                // making two O(log n) searches for the physical index of the slot -- one
                // here and another in Visit(const RunEndEncodedArray&) in case the values
                // is not null.
                if (array_.type()->id() != Type::RUN_END_ENCODED && array_.is_null(index_)) {
                    auto null = MakeNullScalar(array_.type());
                    if (is_dictionary(array_.type()->id())) {
                        auto &dict_null = turbo::checked_cast<DictionaryScalar &>(*null);
                        const auto &dict_array = turbo::checked_cast<const DictionaryArray &>(array_);
                        dict_null.value.dictionary = dict_array.dictionary();
                    }
                    return null;
                }

                TURBO_RETURN_NOT_OK(visit_array_inline(array_, this));
                return std::move(out_);
            }

            ScalarFromArraySlotImpl(const Array &array, int64_t index)
                    : array_(array), index_(index) {}

            const Array &array_;
            int64_t index_;
            std::shared_ptr<Scalar> out_;
        };

    }  // namespace internal

    turbo::Result<std::shared_ptr<Scalar>> Array::get_scalar(int64_t i) const {
        return internal::ScalarFromArraySlotImpl{*this, i}.finish();
    }

    std::string Array::Diff(const Array &other) const {
        std::stringstream diff;
        TURBO_UNUSED(equals(other, EqualOptions().diff_sink(&diff)));
        return diff.str();
    }

    bool Array::equals(const Array &arr, const EqualOptions &opts) const {
        return array_equals(*this, arr, opts);
    }

    bool Array::equals(const std::shared_ptr<Array> &arr, const EqualOptions &opts) const {
        if (!arr) {
            return false;
        }
        return equals(*arr, opts);
    }

    bool Array::approx_equals(const Array &arr, const EqualOptions &opts) const {
        return array_approx_equals(*this, arr, opts);
    }

    bool Array::approx_equals(const std::shared_ptr<Array> &arr,
                             const EqualOptions &opts) const {
        if (!arr) {
            return false;
        }
        return approx_equals(*arr, opts);
    }

    bool Array::range_equals(const Array &other, int64_t start_idx, int64_t end_idx,
                            int64_t other_start_idx, const EqualOptions &opts) const {
        return array_range_equals(*this, other, start_idx, end_idx, other_start_idx, opts);
    }

    bool Array::range_equals(const std::shared_ptr<Array> &other, int64_t start_idx,
                            int64_t end_idx, int64_t other_start_idx,
                            const EqualOptions &opts) const {
        if (!other) {
            return false;
        }
        return array_range_equals(*this, *other, start_idx, end_idx, other_start_idx, opts);
    }

    bool Array::range_equals(int64_t start_idx, int64_t end_idx, int64_t other_start_idx,
                            const Array &other, const EqualOptions &opts) const {
        return array_range_equals(*this, other, start_idx, end_idx, other_start_idx, opts);
    }

    bool Array::range_equals(int64_t start_idx, int64_t end_idx, int64_t other_start_idx,
                            const std::shared_ptr<Array> &other,
                            const EqualOptions &opts) const {
        if (!other) {
            return false;
        }
        return array_range_equals(*this, *other, start_idx, end_idx, other_start_idx, opts);
    }

    std::shared_ptr<Array> Array::slice(int64_t offset, int64_t length) const {
        return make_array(data_->slice(offset, length));
    }

    std::shared_ptr<Array> Array::slice(int64_t offset) const {
        int64_t slice_length = data_->length - offset;
        return slice(offset, slice_length);
    }

    turbo::Result<std::shared_ptr<Array>> Array::slice_safe(int64_t offset, int64_t length) const {
        TURBO_MOVE_OR_RAISE(auto sliced_data, data_->slice_safe(offset, length));
        return make_array(std::move(sliced_data));
    }

    turbo::Result<std::shared_ptr<Array>> Array::slice_safe(int64_t offset) const {
        if (offset < 0) {
            // Avoid UBSAN in subtraction below
            return turbo::out_of_range_error("Negative array slice offset");
        }
        return slice_safe(offset, data_->length - offset);
    }

    std::string Array::to_string() const {
        std::stringstream ss;
        KCHECK_OK(pretty_print(*this, 0, &ss));
        return ss.str();
    }

    void print_to(const Array &x, std::ostream *os) { *os << x.to_string(); }

    turbo::Result<std::shared_ptr<Array>> Array::View(
            const std::shared_ptr<DataType> &out_type) const {
        TURBO_MOVE_OR_RAISE(std::shared_ptr<ArrayData> result,
                              internal::get_array_view(data_, out_type));
        return make_array(result);
    }

    turbo::Result<std::shared_ptr<Array>> Array::copy_to(
            const std::shared_ptr<MemoryManager> &to) const {
        TURBO_MOVE_OR_RAISE(auto copied_data, data()->copy_to(to));
        return make_array(copied_data);
    }

    turbo::Result<std::shared_ptr<Array>> Array::view_or_copy_to(
            const std::shared_ptr<MemoryManager> &to) const {
        TURBO_MOVE_OR_RAISE(auto new_data, data()->view_or_copy_to(to));
        return make_array(new_data);
    }

    // ----------------------------------------------------------------------
    // NullArray

    NullArray::NullArray(int64_t length) {
        set_data(ArrayData::create(null(), length, {nullptr}, length));
    }

    // ----------------------------------------------------------------------
    // Implement Array::Accept as inline visitor

    turbo::Status Array::Accept(ArrayVisitor *visitor) const {
        return visit_array_inline(*this, visitor);
    }

    turbo::Status Array::validate() const { return internal::ValidateArray(*this); }

    turbo::Status Array::validate_full() const { return internal::ValidateArrayFull(*this); }

}  // namespace nebula
